#!/bin/bash

eth_name=$1
total_cpus=$(grep -c processor /proc/cpuinfo)
queues_num=`ethtool -l $eth_name | grep Combined | tail -1 | awk '{print $2}'`

function set_xps_3core_per_queue()
{
    if [ $total_cpus -eq 192 ];then
        echo "00000000,00000000,00000000,00000000,00000000,00000007" > /sys/class/net/$eth_name/queues/tx-0/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000000,00000038" > /sys/class/net/$eth_name/queues/tx-1/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000000,000001c0" > /sys/class/net/$eth_name/queues/tx-2/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000000,00000e00" > /sys/class/net/$eth_name/queues/tx-3/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000000,00007000" > /sys/class/net/$eth_name/queues/tx-4/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000000,00038000" > /sys/class/net/$eth_name/queues/tx-5/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000000,001c0000" > /sys/class/net/$eth_name/queues/tx-6/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000000,00e00000" > /sys/class/net/$eth_name/queues/tx-7/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000000,07000000" > /sys/class/net/$eth_name/queues/tx-8/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000000,38000000" > /sys/class/net/$eth_name/queues/tx-9/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000001,c0000000" > /sys/class/net/$eth_name/queues/tx-10/xps_cpus
        echo "00000000,00000000,00000000,00000000,0000000e,00000000" > /sys/class/net/$eth_name/queues/tx-11/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000070,00000000" > /sys/class/net/$eth_name/queues/tx-12/xps_cpus
        echo "00000000,00000000,00000000,00000000,00000380,00000000" > /sys/class/net/$eth_name/queues/tx-13/xps_cpus
        echo "00000000,00000000,00000000,00000000,00001c00,00000000" > /sys/class/net/$eth_name/queues/tx-14/xps_cpus
        echo "00000000,00000000,00000000,00000000,0000e000,00000000" > /sys/class/net/$eth_name/queues/tx-15/xps_cpus
        echo "00000000,00000000,00000000,00000000,00070000,00000000" > /sys/class/net/$eth_name/queues/tx-16/xps_cpus
        echo "00000000,00000000,00000000,00000000,00380000,00000000" > /sys/class/net/$eth_name/queues/tx-17/xps_cpus
        echo "00000000,00000000,00000000,00000000,01c00000,00000000" > /sys/class/net/$eth_name/queues/tx-18/xps_cpus
        echo "00000000,00000000,00000000,00000000,0e000000,00000000" > /sys/class/net/$eth_name/queues/tx-19/xps_cpus
        echo "00000000,00000000,00000000,00000000,70000000,00000000" > /sys/class/net/$eth_name/queues/tx-20/xps_cpus
        echo "00000000,00000000,00000000,00000003,80000000,00000000" > /sys/class/net/$eth_name/queues/tx-21/xps_cpus
        echo "00000000,00000000,00000000,0000001c,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-22/xps_cpus
        echo "00000000,00000000,00000000,000000e0,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-23/xps_cpus
        echo "00000000,00000000,00000000,00000700,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-24/xps_cpus
        echo "00000000,00000000,00000000,00003800,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-25/xps_cpus
        echo "00000000,00000000,00000000,0001c000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-26/xps_cpus
        echo "00000000,00000000,00000000,000e0000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-27/xps_cpus
        echo "00000000,00000000,00000000,00700000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-28/xps_cpus
        echo "00000000,00000000,00000000,03800000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-29/xps_cpus
        echo "00000000,00000000,00000000,1c000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-30/xps_cpus
        echo "00000000,00000000,00000000,e0000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-31/xps_cpus
        echo "00000000,00000000,00000007,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-32/xps_cpus
        echo "00000000,00000000,00000038,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-33/xps_cpus
        echo "00000000,00000000,000001c0,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-34/xps_cpus
        echo "00000000,00000000,00000e00,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-35/xps_cpus
        echo "00000000,00000000,00007000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-36/xps_cpus
        echo "00000000,00000000,00038000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-37/xps_cpus
        echo "00000000,00000000,001c0000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-38/xps_cpus
        echo "00000000,00000000,00e00000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-39/xps_cpus
        echo "00000000,00000000,07000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-40/xps_cpus
        echo "00000000,00000000,38000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-41/xps_cpus
        echo "00000000,00000001,c0000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-42/xps_cpus
        echo "00000000,0000000e,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-43/xps_cpus
        echo "00000000,00000070,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-44/xps_cpus
        echo "00000000,00000380,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-45/xps_cpus
        echo "00000000,00001c00,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-46/xps_cpus
        echo "00000000,0000e000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-47/xps_cpus
        echo "00000000,00070000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-48/xps_cpus
        echo "00000000,00380000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-49/xps_cpus
        echo "00000000,01c00000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-50/xps_cpus
        echo "00000000,0e000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-51/xps_cpus
        echo "00000000,70000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-52/xps_cpus
        echo "00000003,80000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-53/xps_cpus
        echo "0000001c,00000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-54/xps_cpus
        echo "000000e0,00000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-55/xps_cpus
        echo "00000700,00000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-56/xps_cpus
        echo "00003800,00000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-57/xps_cpus
        echo "0001c000,00000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-58/xps_cpus
        echo "000e0000,00000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-59/xps_cpus
        echo "00700000,00000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-60/xps_cpus
        echo "03800000,00000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-61/xps_cpus
        echo "1c000000,00000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-62/xps_cpus
        echo "e0000000,00000000,00000000,00000000,00000000,00000000" > /sys/class/net/$eth_name/queues/tx-63/xps_cpus
    fi
}

function set_xps()
{
    if [[ $total_cpus -ne 2 && $total_cpus -ne 4 &&  $total_cpus -ne 8 && \
            $total_cpus -ne 16 && $total_cpus -ne 32 && \
            $total_cpus -ne 64 && $total_cpus -ne 128 && \
            $total_cpus -ne 48 && $total_cpus -ne 96 && $total_cpus -ne 192 ]]; then
        echo "unsupport to set xps with cpu number $total_cpus"
        exit 1
    fi
    # 如果队列不是2的整数倍，给取整
    if [[ $queues_num -eq 15 || $queues_num -eq 31 ||  $queues_num -eq 63 ]]; then
        queues_num=$(($queues_num+1))
    fi

    per_queue_cpus=$(($total_cpus/$queues_num))
    pos=32
    i=0
    while [ $i -le $(($queues_num-1)) ]
    do
        case $per_queue_cpus in
        1)    # 每1个CPU 1个队列
            cpu_affinity_1=`echo  "00000000,00000000,00000000,00000000" | sed  "s/0/1/$pos" | grep '[1-z].*' -o`
            cpu_affinity_2=`echo  "00000000,00000000,00000000,00000000" | sed  "s/0/2/$pos" | grep '[1-z].*' -o`
            cpu_affinity_4=`echo  "00000000,00000000,00000000,00000000" | sed  "s/0/4/$pos" | grep '[1-z].*' -o`
            cpu_affinity_8=`echo  "00000000,00000000,00000000,00000000" | sed  "s/0/8/$pos" | grep '[1-z].*' -o`
            echo  $cpu_affinity_1 > /sys/class/net/$eth_name/queues/tx-$i/xps_cpus
            echo  $cpu_affinity_2 > /sys/class/net/$eth_name/queues/tx-$(($i+1))/xps_cpus
            echo  $cpu_affinity_4 > /sys/class/net/$eth_name/queues/tx-$(($i+2))/xps_cpus
            echo  $cpu_affinity_8 > /sys/class/net/$eth_name/queues/tx-$(($i+3))/xps_cpus
            i=$(($i+4))
            ;;
        2)    # 每2个CPU 1个队列
            cpu_affinity_3=`echo  "00000000,00000000,00000000,00000000" | sed  "s/0/3/$pos" | grep '[1-z].*' -o`
            cpu_affinity_c=`echo  "00000000,00000000,00000000,00000000" | sed  "s/0/c/$pos" | grep '[1-z].*' -o`
            echo  $cpu_affinity_3 > /sys/class/net/$eth_name/queues/tx-$i/xps_cpus
            echo  $cpu_affinity_c > /sys/class/net/$eth_name/queues/tx-$(($i+1))/xps_cpus
            i=$(($i+2))
            ;;
        4)    # 每4个CPU 1个队列
            cpu_affinity_f=`echo  "00000000,00000000,00000000,00000000" | sed  "s/0/f/$pos" | grep '[1-z].*' -o`
            echo  $cpu_affinity_f > /sys/class/net/$eth_name/queues/tx-$i/xps_cpus
            i=$(($i+1))
            ;;
        *)
            break
            ;;
        esac
        pos=$(($pos-1))
    done

    if [ $per_queue_cpus -eq 3 ];then    # 每3个CPU 1个队列
        set_xps_3core_per_queue
    fi

    show_xps
}

function unset_xps()
{
   for i in `seq 0 1 $(($queues_num-1))`
   do
       echo 0 > /sys/class/net/$eth_name/queues/tx-$i/xps_cpus
   done
   show_xps
}

function show_xps()
{
   for i in `seq 0 1 $(($queues_num-1))`
   do
       smp_affinity_list=`cat /sys/class/net/$eth_name/queues/tx-$i/xps_cpus`
       echo "/sys/class/net/$eth_name/queues/tx-$i/xps_cpus = $smp_affinity_list"
   done
}


function set_rps()
{
    case $total_cpus in
        2)
            cpu_affinity_rps="f"
            ;;
        4)
            cpu_affinity_rps="f"
            ;;
        8)
            cpu_affinity_rps="ff"
            ;;
        16) 
            cpu_affinity_rps="ffff"
            ;;
        32)
            cpu_affinity_rps="ffffffff"
            ;;
        48)
            cpu_affinity_rps="ffff,ffffffff"
            ;;
        64)
            cpu_affinity_rps="ffffffff,ffffffff"
            ;;
        96)
            cpu_affinity_rps="ffffffff,ffffffff,ffffffff"
            ;;
        128)
            cpu_affinity_rps="ffffffff,ffffffff,ffffffff,ffffffff"
            ;;
        192)
            cpu_affinity_rps="ffffffff,ffffffff,ffffffff,ffffffff,ffffffff,ffffffff"
            ;;
    esac

    for i in `seq 0 1 $(($queues_num-1))`
    do
        echo $cpu_affinity_rps > /sys/class/net/$eth_name/queues/rx-$i/rps_cpus
    done

    rfc=4096
    for i in `seq 0 1 $(($queues_num-1))`
    do
        echo $rfc > /sys/class/net/$eth_name/queues/rx-$i/rps_flow_cnt
    done

    cc=$(grep -c processor /proc/cpuinfo)
    rsfe=$(echo $cc*$rfc | bc)
    sysctl -w net.core.rps_sock_flow_entries=$rsfe   #default 0

    show_rps
}

function unset_rps()
{
    for i in `seq 0 1 $(($queues_num-1))`
    do
        echo 0 > /sys/class/net/$eth_name/queues/rx-$i/rps_cpus
    done

    for i in `seq 0 1 $(($queues_num-1))`
    do
        echo 0 > /sys/class/net/$eth_name/queues/rx-$i/rps_flow_cnt
    done
    sysctl -w net.core.rps_sock_flow_entries=0
    show_rps
}

function show_rps()
{
    for i in `seq 0 1 $(($queues_num-1))`
    do
        smp_affinity_list=`cat /sys/class/net/$eth_name/queues/rx-$i/rps_cpus`
        echo "/sys/class/net/$eth_name/queues/rx-$i/rps_cpus = $smp_affinity_list"
    done

    for i in `seq 0 1 $(($queues_num-1))`
    do
        rps_flow_cnt=`cat /sys/class/net/$eth_name/queues/rx-$i/rps_flow_cnt`
        echo "/sys/class/net/$eth_name/queues/rx-$i/rps_flow_cnt = $rps_flow_cnt"
    done

    sysctl net.core.rps_sock_flow_entries
}

function usage()
{
   echo "Usage: 
      sh $0 eth_name [ xps | rps ] [ show | enable | disable ]"
}

if [[ "$#" -ne 3 && "$2" != "xps" && "$2" != "rps" ]]; then
  usage
  exit 1
fi

case $3 in
  show)
    echo "Executing command show $2"
    if [ $2 == "xps" ];then show_xps; else show_rps; fi
    ;;
  enable)
    echo "Executing command set $2"
    if [ $2 == "xps" ];then set_xps; else set_rps; fi
    ;;
  disable)
    echo "Executing command unset $2"
    if [ $2 == "xps" ];then unset_xps; else unset_rps; fi
    ;;
  *)
    echo "Invalid command: $3"
    usage
    exit 1
    ;;
esac